from core.function import Function
from core.cuda import cuda
from utils.common import to_array
import utils

class Sub(Function):
    def __init__(self):
        self.x1_shape = None
        self.x0_shape = None

    def forward(self, x0, x1):
        """
        执行 Sub 操作的前向传播。

        Parameters:
            x0 (Variable): 第一个输入变量。
            x1 (Variable): 第二个输入变量。

        Returns:
            Variable: Sub 操作的结果。
        """
        self.x0_shape, self.x1_shape = x0.shape, x1.shape
        x0, x1 = cuda.to_array(x0), cuda.to_array(x1)
        y = x0 - x1
        return y

    def backward(self, gy):

        """
        执行 Sub 操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """
        gx0 = gy
        gx1 = -gy
        if self.x0_shape != self.x1_shape:  # for broadcast
            gx0 = utils.functions_collect.sum_to(gx0, self.x0_shape)
            gx1 = utils.functions_collect.sum_to(gx1, self.x1_shape)

        return gx0, gx1

    def sub(self, x0, x1):
        """
        执行 Sub 操作。

        Parameters:
            x0 (Variable): 第一个输入变量。
            x1 (Variable): 第二个输入变量。

        Returns:
            Variable: Sub 操作的结果。
        """
        x1 = to_array(x1)
        return self(x0, x1)
